import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import cv2  # OpenCV 用于摄像头捕捉


class MonNet(nn.Module):
    def __init__(self):
        super(MonNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, kernel_size=(5, 5), padding=0)
        self.bn1 = nn.BatchNorm2d(6)
        self.maxpool1 = nn.MaxPool2d(6)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(6, 16, kernel_size=(5, 5), padding=5)
        self.bn1 = nn.BatchNorm2d(16)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 5)
        self.sf1 = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.maxpool1(x)
        x = self.relu(self.conv2(x))
        x = self.maxpool1(x)
        x = x.view(-1, 16 * 5 * 5)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        y = self.sf1(self.fc3(x))
        return y


# 加载检查点
model = MonNet()
checkpoint = torch.load('checkpoint.pth', map_location=torch.device('cpu'))  # 如果使用CPU加载模型
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()  # 切换到评估模式

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((160, 160)),
    transforms.ToTensor(),
])


# 预测函数
def predict_image(image):
    # 预处理图像
    image = transform(image).unsqueeze(0)  # 增加batch维度
    with torch.no_grad():  # 禁用梯度计算
        output = model(image)
        _, predicted = torch.max(output, 1)
    return predicted.item()


# 对摄像头捕获的图像进行预测
def predict_from_camera():
    cap = cv2.VideoCapture(0)  # 打开摄像头
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # 显示摄像头画面
        cv2.imshow('Press "q" to capture and classify', frame)

        # 按下 "q" 键拍照并预测
        if cv2.waitKey(1) & 0xFF == ord('q'):
            # 将图像从BGR转换为RGB格式
            image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            image = Image.fromarray(image)  # 转换为 PIL 图像格式

            # 调用预测函数
            class_idx = predict_image(image)
            print(f"Predicted Class from Camera: {class_idx}")
            break

    # 释放摄像头资源
    cap.release()
    cv2.destroyAllWindows()


# 对特定图像文件进行预测
def predict_from_file(image_path):
    image = Image.open(image_path).convert('RGB')  # 打开并转换为 RGB
    class_idx = predict_image(image)
    print(f"Predicted Class from Image File: {class_idx}")


# # 使用摄像头进行预测
# predict_from_camera()

#或者对特定图像文件进行预测
predict_from_file(r'E:\Desktop\实验六\DA_Monkeys\White face saki\WFS (13)_aug_0.jpg')
